import csv
from clickhouse_driver import Client
from datetime import datetime, timedelta
import os

# 获取脚本所在路径
curPath = os.path.split(os.path.realpath(__file__))[0]


def date_sub(day):
    today = datetime.now()
    three_days_ago = today - timedelta(days=day)
    # 将时间格式化为YYYY-MM-DD格式
    date_str = three_days_ago.strftime('%Y-%m-%d')
    return date_str


def find_index(lst, target):
    for i, item in enumerate(lst):
        if item == target:
            return i
    return -1


def col_type_convert(schema, origin_data):
    # 将数据转换为指定类型
    for row in origin_data:
        for idx, field in enumerate(schema):
            if field[1].lower().find('int') != -1:
                row[idx] = int(row[idx])


def write2ch(client, table, datas, time_fields):
    # 获取表中的字段名称
    schema = client.execute("DESCRIBE TABLE {}".format(table))
    fields = [field[0] for field in schema]
    insert_sql = "INSERT INTO {} ({}) VALUES".format(table, ','.join(fields))
    if len(datas) > 0:
        # 将数据中的时间改为当前时间
        for row in datas:
            for time_field in time_fields:
                row[find_index(fields, time_field)] = datetime.now()

        print(insert_sql)
        client.execute(insert_sql, datas, types_check=False)
    else:
        # 从CSV文件中读取数据
        with open('{}/data/{}.csv'.format(curPath, table), 'r', newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            csv_data = [row for row in reader]
            col_type_convert(schema, csv_data)
            write2ch(client, table, csv_data, time_fields)


def read_ch(client, table, time_field):
    # 创建一个查询，用于从表中读取数据
    filter_time = date_sub(1000)
    query = "SELECT * FROM {} WHERE {} > '{}' limit 50".format(table, time_field, filter_time)
    print(query)

    results = client.execute(query)

    # 将数据写入文件
    if len(results) > 0:
        with open('{}/{}.csv'.format(curPath, table), 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            for row in results:
                writer.writerow(row)
    # 将结果从元组转换为列表
    return [list(sublist) for sublist in results]


if __name__ == '__main__':
    # 创建 ch 客户端连接
    # src = Client(host='10.10.22.36', port=9100, user='default', password='d123456', database='network_security')
    dst_27 = Client(host='10.10.3.27', port=9200, user='default', password='d123456', database='network_security')
    dst_30 = Client(host='10.10.3.30', port=9100, user='default', password='d123456', database='network_security')

    # 定义要导出数据的表
    ch_table_dict = {
        'td_host_control_event': ['found_time'],
        'td_network_attack_event': ['found_time'],
        'td_harm_program_event': ['found_time'],
        'td_special_topic_event': ['event_time'],
        'td_evil_packet_monitor_resp': ['event_time', 'create_time'],
        'td_evil_file_monitor_resp': ['access_time', 'create_time'],
        'td_evil_resource_monitor_resp': ['event_time', 'create_time'],
        'td_deal_commond_resp': ['gather_time', 'create_time'],
        'td_xdr_log_http': ['start_time'],
        'td_xdr_log_dns': ['start_time'],
        'td_xdr_log_ftp': ['start_time'],
        'td_xdr_log_email': ['start_time'],
        'td_stock_resp': ['report_time', 'create_time'],
        # 'td_device_status_detail': ['create_time'], 报错未解决
        'tv_his_minitor_resp': ['create_time', 'event_time'],
        'tv_his_evil_sample_resp': ['create_time', 'access_time']
    }

    print('insert 10.10.3.27')
    for key, val in ch_table_dict.items():
        # data = read_ch(src_client, key, val[0])
        data = []
        write2ch(dst_27, key, data, val)

    print('insert 10.10.3.30')
    for key, val in ch_table_dict.items():
        # data = read_ch(src_client, key, val[0])
        data = []
        write2ch(dst_30, key, data, val)
